# 运行过程中要对原始图片鼠标左键点选四个角点，顺序必须为左上、右上、左下，右下
# https://blog.csdn.net/weixin_42326479/article/details/102938733
import cv2
import os
import numpy as np
import tkinter as tk
import tkinter.ttk as ttk
import copy
from tkinter import messagebox
from tkinter import filedialog

np.set_printoptions(suppress=True)

root = tk.Tk()
root.geometry("500x500")
root.title("图片矫正处理工具")

# from tqdm import tqdm

# arr = [[0,0],[img.shape[1],0],[0,img.shape[0]],[img.shape[1],img.shape[0]]]
arr = []
# rc_FILE = "D:/data/IMG_20230719_105053.jpg"
img = None
image2 = None
WIN_NAME = "top_set"
top_type = "near"


class Rect(object):
    def __init__(self):
        self.tl = (0, 0)
        self.br = (0, 0)

    def regularize(self):
        """
        make sure tl = TopLeft point, br = BottomRight point
        """
        pt1 = (min(self.tl[0], self.br[0]), min(self.tl[1], self.br[1]))
        pt2 = (max(self.tl[0], self.br[0]), max(self.tl[1], self.br[1]))
        self.tl = pt1
        self.br = pt2

    def get_points(self):
        return [list(self.tl), [self.tl[0], self.br[1]], [self.br[0], self.tl[1]], list(self.br)]


class DrawRects(object):
    def __init__(self, image, color, thickness=1):
        self.original_image = image
        self.image_for_show = image.copy()
        self.color = color
        self.thickness = thickness
        self.rects = []
        self.current_rect = Rect()
        self.left_button_down = False

    @staticmethod
    def __clip(value, low, high):
        """
        clip value between low and high
        Parameters
        ----------
        value: a number
            value to be clipped
        low: a number
            low limit
        high: a number
            high limit
        Returns
        -------
        output: a number
            clipped value
        """
        output = max(value, low)
        output = min(output, high)
        return output

    def shrink_point(self, x, y):
        """
        shrink point (x, y) to inside image_for_show
        Parameters
        ----------
        x, y: int, int
            coordinate of a point
        Returns
        -------
        x_shrink, y_shrink: int, int
            shrinked coordinate
        """
        height, width = self.image_for_show.shape[0:2]
        x_shrink = self.__clip(x, 0, width)
        y_shrink = self.__clip(y, 0, height)
        return (x_shrink, y_shrink)

    def append(self):
        """
        add a rect to rects list
        """
        self.rects.append(copy.deepcopy(self.current_rect))

    def pop(self):
        """
        pop a rect from rects list
        """
        rect = Rect()
        if self.rects:
            rect = self.rects.pop()
        return rect

    def reset_image(self):
        """
        reset image_for_show using original image
        """
        self.image_for_show = self.original_image.copy()

    def draw_all_rects(self):
        """
        draw rects on image_for_show
        """
        for rect in self.rects:
            cv2.rectangle(self.image_for_show, rect.tl, rect.br,
                          color=self.color, thickness=self.thickness, lineType=cv2.LINE_AA)

    def draw_current_rect(self):
        """
        draw current rect on image_for_show
        """
        cv2.rectangle(self.image_for_show,
                      self.current_rect.tl, self.current_rect.br,
                      color=self.color, thickness=self.thickness)


def onmouse_draw_rect(event, x, y, flags, draw_rects):
    if event == cv2.EVENT_LBUTTONDOWN:
        # pick first point of rect
        print('pt1: x = %d, y = %d' % (x, y))
        draw_rects.left_button_down = True
        draw_rects.current_rect.tl = (x, y)
    if draw_rects.left_button_down and event == cv2.EVENT_MOUSEMOVE:
        # pick second point of rect and draw current rect
        draw_rects.current_rect.br = draw_rects.shrink_point(x, y)
        # draw_rects.reset_image()
        # draw_rects.draw_all_rects()
        # draw_rects.draw_current_rect()
    if event == cv2.EVENT_LBUTTONUP:
        # finish drawing current rect and append it to rects list
        draw_rects.left_button_down = False
        draw_rects.current_rect.br = draw_rects.shrink_point(x, y)
        print('pt2: x = %d, y = %d' % (draw_rects.current_rect.br[0],
                                       draw_rects.current_rect.br[1]))
        draw_rects.current_rect.regularize()
        draw_rects.append()
        draw_rects.reset_image()
        # draw_rects.draw_all_rects()
        draw_rects.draw_current_rect()
        global WIN_NAME
        cv2.imshow(WIN_NAME, draw_rects.image_for_show)
        # cv2.imwrite("D:/data/test.jpg", draw_rects.image_for_show)
    if (not draw_rects.left_button_down) and event == cv2.EVENT_RBUTTONDOWN:
        # pop the last rect in rects list
        draw_rects.pop()
        draw_rects.reset_image()
        # draw_rects.draw_all_rects()
        # draw_rects.draw_current_rect()


def on_EVENT_LBUTTONDOWN(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        arr.append([x, y])
        cv2.circle(img, (x, y), 8, (255, 0, 255), thickness=3)
        cv2.imshow("manual", img)
        if len(arr) == 4:
            # 四边形的四个顶点
            pts = np.array([arr[0], arr[1], arr[3], arr[2]], np.int32)
            cv2.polylines(img, [pts], isClosed=True, color=(0, 255, 0), thickness=2)
            cv2.imshow("manual", img)


def doImg(src_FILE):
    global arr
    arr = []
    global img
    img = cv2.imread(src_FILE)  # 中文目录会读不到
    # img = cv2.imdecode(np.fromfile(src_FILE, dtype=np.uint8), -1)
    img_transform = img.copy()

    cv2.namedWindow("manual", cv2.WINDOW_FREERATIO)
    cv2.setMouseCallback("manual", on_EVENT_LBUTTONDOWN)
    cv2.imshow("manual", img)
    # Esc 销毁窗口
    if len(arr) == 4 or cv2.waitKey(0) == 27:
        cv2.destroyAllWindows()
        if len(arr) == 3:
            print("选定的三个角点为：{}{}{}".format(arr[0], arr[1], arr[2]))
            return
        elif len(arr) == 2:
            print("选定的两个角点为：{}{}".format(arr[0], arr[1]))
            return
        elif len(arr) == 1:
            print("选定的一个角点为：{}".format(arr[0]))
            return
    # if len(arr) == 2:   #根据平行关系，自动补充p3，p4
    #     arr.append([arr[0][0], img.shape[0] - arr[0][1]])  # (p1.x,h-p2.y)
    #     arr.append([arr[0][1], img.shape[0]])  # (p2.x,h)
    P1, P2, P3, P4 = arr[:]

    print("选定的四个角点为：{}{}{}{}".format(P1, P2, P3, P4))  # Z字型选中四点
    cv2.circle(img_transform, tuple(P1), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P2), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P3), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P4), 2, (255, 0, 0), 3)
    # # 输入梯形的四个顶点
    srcPoints = np.vstack((P1, P2, P3, P4))
    srcPoints = np.float32(srcPoints)
    # 目标的像素值大小
    long = img.shape[1]  # 图片的长宽
    short = img.shape[0]
    # 设置目标画布的大小
    canvasPoints = np.array([[0, 0], [int(long), 0], [0, int(short)], [int(long), int(short)]])
    canvasPoints = np.float32(canvasPoints)
    # 计算转换矩阵
    perspectiveMatrix = cv2.getPerspectiveTransform(srcPoints, canvasPoints)
    print(perspectiveMatrix)
    # 完成透视变换
    perspectiveImg = cv2.warpPerspective(img_transform, perspectiveMatrix, (int(long), int(short)))
    print("展示校正结果")
    file_base_name = os.path.basename(src_FILE)
    cv2.namedWindow(file_base_name, cv2.WINDOW_FREERATIO)
    cv2.imshow(file_base_name, perspectiveImg)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    dir_out = os.path.join(os.path.dirname(src_FILE), "cr")
    if not os.path.exists(dir_out):
        os.mkdir(dir_out)
    dst_FILE = os.path.join(dir_out, file_base_name)  # "D:/data/real/out/result.JPG"
    cv2.imwrite(dst_FILE, perspectiveImg)  # 中文目录写入不成功，需要手动保存
    print("保存结果完毕到"+dst_FILE)
    if file_base_name.startswith("FAR") or file_base_name.startswith("NEAR"):
        setting_filepath = "C:\\uni-orange\\"
        if file_base_name.startswith("FAR"):
            setting_filepath = setting_filepath + "far.npy"
        elif file_base_name.startswith("NEAR"):
            setting_filepath = setting_filepath + "near.npy"
        is_save = messagebox.askyesno("提示", "是否保存到配置文件"+setting_filepath)
        if is_save:
            point_array = np.array(arr)
            np.save(setting_filepath, point_array)
            messagebox.showinfo("提示", "成功保存到 "+setting_filepath)

def choose_image():
    src_FILE = filedialog.askopenfilename(initialdir="images/", title="选择图片",
                                          filetypes=(("JPEG files", "*.jpg"), ("PNG files", "*.png")))
    print(src_FILE)
    doImg(src_FILE)


def choose_image_far():
    src_FILE = filedialog.askopenfilename(initialdir="images/", title="选择顶拍远侧图片",
                                          filetypes=(("JPEG files", "*.jpg"), ("PNG files", "*.png")))
    print(src_FILE)
    global top_type
    top_type = "far"
    global image2
    image2 = cv2.imread(src_FILE)
    draw_rects = DrawRects(image2, (0, 255, 0), 5)
    cv2.namedWindow(WIN_NAME, 0)
    cv2.imshow(WIN_NAME, draw_rects.image_for_show)
    cv2.setMouseCallback(WIN_NAME, onmouse_draw_rect, draw_rects)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    setting_filepath = "C:\\uni-orange\\far.npy"
    is_save = messagebox.askyesno("提示", "是否保存到配置文件"+setting_filepath)
    if is_save:
        point_array = np.array(draw_rects.current_rect.get_points())
        np.save(setting_filepath, point_array)
        messagebox.showinfo("提示", "成功保存到 "+setting_filepath)


def choose_image_near():
    src_FILE = filedialog.askopenfilename(initialdir="images/", title="选择顶拍近侧图片",
                                          filetypes=(("JPEG files", "*.jpg"), ("PNG files", "*.png")))
    print(src_FILE)
    global top_type
    top_type = "near"
    global image2
    image2 = cv2.imread(src_FILE)
    draw_rects = DrawRects(image2, (0, 255, 0), 5)
    cv2.namedWindow(WIN_NAME, 0)
    cv2.imshow(WIN_NAME, draw_rects.image_for_show)
    cv2.setMouseCallback(WIN_NAME, onmouse_draw_rect, draw_rects)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    setting_filepath = "C:\\uni-orange\\near.npy"
    is_save = messagebox.askyesno("提示", "是否保存到配置文件"+setting_filepath)
    if is_save:
        point_array = np.array(draw_rects.current_rect.get_points())
        np.save(setting_filepath, point_array)
        messagebox.showinfo("提示", "成功保存到 "+setting_filepath)


tab_container = ttk.Notebook(root)
tab_container.pack(fill='both', expand=1)

frame1 = tk.Frame(tab_container)
tab_container.add(frame1, text="Z字形选中四点矫正")
frame2 = tk.Frame(tab_container)
tab_container.add(frame2, text="顶拍矩形矫正设置")

button = tk.Button(frame1, text="选择图片", command=choose_image)
button.pack(side="top", pady=50)

button22 = tk.Button(frame2, text="选择顶拍远侧照片", command=choose_image_far)
button22.pack(side="top", pady=50)
button21 = tk.Button(frame2, text="选择顶拍近侧照片", command=choose_image_near)
button21.pack(side="top", pady=50)

root.mainloop()
